#include "PlanePostProcessMatcher.h"
#include <vector>

using namespace std;

PlanePostProcessMatcher::PlanePostProcessMatcher()
{
	name = "PlanePostProcessMatcher";
}

PlanePostProcessMatcher::PlanePostProcessMatcher(FrameMatcher * base_)
{
	name = "PlanePostProcessMatcher";
	base = base_;
}

PlanePostProcessMatcher::~PlanePostProcessMatcher(){printf("delete PlanePostProcessMatcher\n");}

Eigen::Matrix4d PlanePostProcessMatcher::optimize(vector<segmentation_fit *> match_src, vector<segmentation_fit *> match_dst,Transformation * trans){
	Eigen::Matrix4d transformationMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i< 4; i++){for(int j = 0; j < 4; j++){transformationMat(i,j) = trans->transformationMatrix(i,j);}}
	
	//cout<<"aa:\n"<<aa<<"\n";
	//double improvement = 1;
	for(int i = 0; i < 1000; i++)
	{
		//printf("iter: %i\n",i);
		double current_error = getError(transformationMat,match_src,match_dst,trans);
		double * current_gradient = getGradient(transformationMat,match_src,match_dst,trans);
		double gradient_length = sqrt(current_gradient[0]*current_gradient[0] + current_gradient[1]*current_gradient[1] + current_gradient[2]*current_gradient[2] + current_gradient[3]*current_gradient[3] + current_gradient[4]*current_gradient[4] + current_gradient[5]*current_gradient[5]);
		double gradient_mul = 10.0*current_error/gradient_length;
		current_gradient[0]*=gradient_mul;
		current_gradient[1]*=gradient_mul;
		current_gradient[2]*=gradient_mul;
		current_gradient[3]*=gradient_mul;
		current_gradient[4]*=gradient_mul;
		current_gradient[5]*=gradient_mul;
		//printf("current_error: %.15f\n",current_error);
		//printf("gradient_length: %.15f\n",gradient_length);
		//printf("gradient_mul: %.5f\n",gradient_mul);
		transformationMat = useGradient(transformationMat,current_gradient);
	}
	for(int i = 0; i< 4; i++){for(int j = 0; j < 4; j++){trans->transformationMatrix(i,j) = transformationMat(i,j);}}
	return transformationMat;
}

Eigen::Matrix4d PlanePostProcessMatcher::useGradient(Eigen::Matrix4d transformationMat,double * gradient){
	//printf("useGradient\n");

	
	Eigen::Matrix4d tmpMat;
	Eigen::Matrix3d rot;
	
	rot = Eigen::AngleAxisd(2.0*0.001*-gradient[0], Eigen::Vector3d(1, 0, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){for(int j = 0; j<3; j++){tmpMat(i,j) = rot(i,j);}}
	transformationMat=transformationMat*tmpMat;
	
	rot = Eigen::AngleAxisd(2.0*0.001*-gradient[1], Eigen::Vector3d(0, 1, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){for(int j = 0; j<3; j++){tmpMat(i,j) = rot(i,j);}}
	transformationMat=transformationMat*tmpMat;
	
	rot = Eigen::AngleAxisd(2.0*0.001*-gradient[2], Eigen::Vector3d(0, 0, 1));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){for(int j = 0; j<3; j++){tmpMat(i,j) = rot(i,j);}}
	transformationMat=transformationMat*tmpMat;
	
	return transformationMat;
}

double * PlanePostProcessMatcher::getGradient(Eigen::Matrix4d transformationMat,vector<segmentation_fit *> match_src, vector<segmentation_fit *> match_dst,Transformation * trans){
	//printf("getGradient\n");
	Eigen::Matrix4d tmpMat;
	double * gradient = new double[6];

	Eigen::Matrix3d pxrot;
	pxrot = Eigen::AngleAxisd(0.001, Eigen::Vector3d(1, 0, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = pxrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errorpx = getError(tmpMat,match_src,match_dst,trans);
	
	Eigen::Matrix3d nxrot;
	nxrot = Eigen::AngleAxisd(-0.001,Eigen::Vector3d(1, 0, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = nxrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errornx = getError(tmpMat,match_src,match_dst,trans);
	
	Eigen::Matrix3d pyrot;
	pyrot = Eigen::AngleAxisd(0.001, Eigen::Vector3d(0, 1, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = pyrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errorpy = getError(tmpMat,match_src,match_dst,trans);
	
	Eigen::Matrix3d nyrot;
	nyrot = Eigen::AngleAxisd(-0.001,Eigen::Vector3d(0, 1, 0));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = nyrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errorny = getError(tmpMat,match_src,match_dst,trans);
	
	Eigen::Matrix3d pzrot;
	pzrot = Eigen::AngleAxisd(0.001, Eigen::Vector3d(0, 0, 1));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = pzrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errorpz = getError(tmpMat,match_src,match_dst,trans);
	
	Eigen::Matrix3d nzrot;
	nzrot = Eigen::AngleAxisd(-0.001, Eigen::Vector3d(0, 0, 1));
	tmpMat = Eigen::Matrix4d::Identity();
	for(int i = 0; i < 3; i++){
		for(int j = 0; j<3; j++){
			tmpMat(i,j) = nzrot(i,j);
		}
	}
	tmpMat = transformationMat*tmpMat;
	double errornz = getError(tmpMat,match_src,match_dst,trans);
	//printf("[%f %f %f %f %f %f] -> [%f %f %f]\n",errorpx,errornx,errorpy,errorny,errorpz,errornz,errorpx-errornx,errorpy-errorny,errorpz-errornz);
	gradient[0] = errorpx-errornx;
	gradient[1] = errorpy-errorny;
	gradient[2] = errorpz-errornz;
	gradient[3] = 0;
	gradient[4] = 0;
	gradient[5] = 0;
	return gradient;
}

double PlanePostProcessMatcher::getError(Eigen::Matrix4d transformationMat,vector<segmentation_fit *> match_src, vector<segmentation_fit *> match_dst,Transformation * trans){
	double mat00 = transformationMat(0,0);
	double mat01 = transformationMat(0,1);
	double mat02 = transformationMat(0,2);
	double mat10 = transformationMat(1,0);
	double mat11 = transformationMat(1,1);
	double mat12 = transformationMat(1,2);
	double mat20 = transformationMat(2,0);
	double mat21 = transformationMat(2,1);
	double mat22 = transformationMat(2,2);
	
	double rdiff = 0;
	for(int i = 0; i < 3; i++){
		for(int j = 0; j < 3; j++){
			rdiff+=fabs(transformationMat(i,j)-trans->transformationMatrix(i,j));
		}
	}
	//printf("------------------------------------------\n");

	double angle_error = 0;
	for(int i = 0; i < match_src.size(); i++){
		segmentation_fit * src_sf = match_src.at(i);
		segmentation_fit * dst_sf = match_dst.at(i);
		double src_normal_x =src_sf->U(0,2);
		double src_normal_y =src_sf->U(1,2);
		double src_normal_z =src_sf->U(2,2);
		
		double dst_normal_x =dst_sf->U(0,2);
		double dst_normal_y =dst_sf->U(1,2);
		double dst_normal_z =dst_sf->U(2,2);
		
		double tmp_x = src_normal_x*mat00+src_normal_y*mat01+src_normal_z*mat02;
		double tmp_y = src_normal_x*mat10+src_normal_y*mat11+src_normal_z*mat12;
		double tmp_z = src_normal_x*mat20+src_normal_y*mat21+src_normal_z*mat22;
		
		double angle = tmp_x*dst_normal_x + tmp_y*dst_normal_y + tmp_z*dst_normal_z;
		angle_error+=1-fabs(angle);
	}
	return angle_error;
}
void PlanePostProcessMatcher::update(){}

using namespace std;
const bool debugg_PlanePostProcessMatcher = false;

Transformation * PlanePostProcessMatcher::getTransformation(RGBDFrame * src, RGBDFrame * dst)
{	
	Transformation * base_transformation = base->getTransformation(src, dst);
	int segments_src = src->segments->size()+1;
	int segments_dst = dst->segments->size()+1;
	int ** overlap = new int*[segments_src];
	for(int i = 0; i < segments_src; i++)
	{
		overlap[i] = new int[segments_dst];
		for(int j = 0; j < segments_dst; j++)
		{
			overlap[i][j] = 0;
		}
	}
	
	for(int i = 0; i < 640; i++)
	{
		for(int j = 0; j < 480; j++)
		{
			overlap[src->segmentation[i][j]][dst->segmentation[i][j]]++;
		}
	}
	
	printf("overlap:\n");
	for(int i = 0; i < segments_src; i++)
	{
		for(int j = 0; j < segments_dst; j++)
		{
			printf("%7i ", overlap[i][j]);
		}
		printf("\n");
	}
	
	IplImage* img_combine;
	int width;
	int height;
	if(debugg_PlanePostProcessMatcher)
	{	
		IplImage* rgb_img_src 	= cvLoadImage(src->input->rgb_path.c_str(),CV_LOAD_IMAGE_UNCHANGED);
		char * data_src = (char *)rgb_img_src->imageData;
		IplImage* rgb_img_dst 	= cvLoadImage(dst->input->rgb_path.c_str(),CV_LOAD_IMAGE_UNCHANGED);
		char * data_dst = (char *)rgb_img_dst->imageData;
		
		width = rgb_img_src->width;
		height = rgb_img_src->height;
		
		img_combine = cvCreateImage(cvSize(2*width,height), IPL_DEPTH_8U, 3);
		char * data = (char *)img_combine->imageData;
		


		int index = 0;
		
		for (int j = 0; j < height; j++)
		{
			for (int i = 0; i < width; i++)
			{
				
				int ind = 3*(640*j+i);
				data[3 * (j * (2*width) + (width+i)) + 0] = data_dst[ind +0];
				data[3 * (j * (2*width) + (width+i)) + 1] = data_dst[ind +1];
				data[3 * (j * (2*width) + (width+i)) + 2] = data_dst[ind +2];

				data[3 * (j * (2*width) + (i)) + 0] = data_src[ind +0];
				data[3 * (j * (2*width) + (i)) + 1] = data_src[ind +1];
				data[3 * (j * (2*width) + (i)) + 2] = data_src[ind +2];
				
			}
		}
		
		
		for(int i = 0; i < src->segments->size(); i++)
		{
			segmentation_fit * seg = src->segments->at(i);
			vector< int > * seg_w = seg->seg_w;
			vector< int > * seg_h = seg->seg_h;
			for(int j = 0; j < seg_w->size(); j++){
				unsigned int w = seg_w->at(j);
				unsigned int h = seg_h->at(j);
				data[3 * (h * (2*width) + (w)) + 0] = 255;
				data[3 * (h * (2*width) + (w)) + 1] = 0;
				data[3 * (h * (2*width) + (w)) + 2] = 255;
			}
		}
		
		for(int i = 0; i < dst->segments->size(); i++)
		{
			segmentation_fit * seg = dst->segments->at(i);
			vector< int > * seg_w = seg->seg_w;
			vector< int > * seg_h = seg->seg_h;
			for(int j = 0; j < seg_w->size(); j++){
				unsigned int w = seg_w->at(j);
				unsigned int h = seg_h->at(j);
				data[3 * (h * (2*width) + (width+w)) + 0] = 255;
				data[3 * (h * (2*width) + (width+w)) + 1] = 0;
				data[3 * (h * (2*width) + (width+w)) + 2] = 255;
			}
		}
		cvNamedWindow("src image", CV_WINDOW_AUTOSIZE );
		cvShowImage("src image", rgb_img_src);
		
		cvNamedWindow("dst image", CV_WINDOW_AUTOSIZE );
		cvShowImage("dst image", rgb_img_dst);
		
		cvNamedWindow("combined image", CV_WINDOW_AUTOSIZE );
		cvShowImage("combined image", img_combine);
		
		cvWaitKey(0);
		cvReleaseImage( &rgb_img_src );
		cvReleaseImage( &rgb_img_dst );
	}
	vector<segmentation_fit *> match_src;
	vector<segmentation_fit *> match_dst;
	for(int i = 1; i < segments_src; i++)
	{
		for(int j = 1; j < segments_dst; j++)
		{
			if(overlap[i][j] > 0)
			{
				float p_dst_given_src = float(overlap[i][j])/float(src->segments->at(i-1)->seg_w->size());
				float p_src_given_dst = float(overlap[i][j])/float(dst->segments->at(j-1)->seg_w->size());
				if(p_src_given_dst > 0.25f && p_dst_given_src > 0.25f && (p_src_given_dst + p_dst_given_src) > 0.75f){
					match_src.push_back(src->segments->at(i-1));
					match_dst.push_back(dst->segments->at(j-1));
				}
				if(debugg_PlanePostProcessMatcher){
					//printf("p(%i_dst|%i_src) = %.5f\n",j-1,i-1, p_dst_given_src);
					//printf("p(%i_src|%i_dst) = %.5f\n",i-1,j-1, p_src_given_dst);
					float src_w_mid = 0;
					float src_h_mid = 0;
					for(int k = 0; k < src->segments->at(i-1)->seg_w->size(); k++)
					{
						src_w_mid+= src->segments->at(i-1)->seg_w->at(k);
						src_h_mid+= src->segments->at(i-1)->seg_h->at(k);
					}
					src_w_mid/=float(src->segments->at(i-1)->seg_w->size());
					src_h_mid/=float(src->segments->at(i-1)->seg_w->size());
					
					float dst_w_mid = 0;
					float dst_h_mid = 0;
					for(int k = 0; k < dst->segments->at(j-1)->seg_w->size(); k++)
					{
						dst_w_mid+= dst->segments->at(j-1)->seg_w->at(k);
						dst_h_mid+= dst->segments->at(j-1)->seg_h->at(k);
					}
					dst_w_mid/=float(dst->segments->at(j-1)->seg_w->size());
					dst_h_mid/=float(dst->segments->at(j-1)->seg_w->size());
					if(p_src_given_dst > 0.25f && p_dst_given_src > 0.25f && (p_src_given_dst + p_dst_given_src) > 0.75f){
						cvLine(img_combine,cvPoint(dst_w_mid  + width ,dst_h_mid),cvPoint(src_w_mid,src_h_mid),cvScalar(0, 255, 0, 0),1, 8, 0);
					}
				}
			}
		}
		if(debugg_PlanePostProcessMatcher){printf("\n");}
	}
	optimize(match_src, match_dst,base_transformation);
	
	for(int i = 0; i < segments_src; i++)
	{
		delete[] overlap[i];
	}
	delete[] overlap;
	if(debugg_PlanePostProcessMatcher){
		cvNamedWindow("combined image", CV_WINDOW_AUTOSIZE );
		cvShowImage("combined image", img_combine);
		
		cvWaitKey(0);
		printf("done\n");
		cvReleaseImage( &img_combine );
	}

	return base_transformation;
}
